#include <stdio.h>
#include <unistd.h>
#include <string.h>
#include <stdlib.h>
#include <errno.h>
#include <sys/resource.h>
#include <sys/times.h>
#include <sys/time.h>
#include <fcntl.h>

#include "utils.h"

const char* dbg_names[] = {"STRT", "MGRT", "TASK", "BALN", "SMRY", "NUMA", "CORE"};
const char* dbg_colors[] = {_BLUE_, _GREEN_, _CYAN_, _MAGENTA_, _GREEN_, _BROWN_, 
                            _LIGHTBLUE_};

void check_sched_compat_yield() {
	// first check to see if this version supports CFS, otherwise yield is the old POSIX compliant 
	// version
	FILE* vfile;
	char* vfile_name = "/proc/version";
	if ((vfile = fopen(vfile_name, "r")) == NULL) {
		FAIL("%s: %s", vfile_name, strerror(errno));
	}
	int maj_ver, min_ver;
	if (fscanf(vfile, "Linux version 2.%d.%d", &maj_ver, &min_ver) == EOF) {
		FAIL("fscanf on %s: %s", vfile_name, strerror(errno));
	}
	fclose(vfile);
	if (maj_ver < 6 || min_ver < 23) return;

	FILE* pfile;
	char* pfile_name = "/proc/sys/kernel/sched_compat_yield"; 
	if ((pfile = fopen(pfile_name, "r")) == NULL) {
		FAIL("%s: %s", pfile_name, strerror(errno));
	}
	int compat;
	if (fscanf(pfile, "%d", &compat) == EOF) {
		FAIL("fscanf on %s: %s", pfile_name, strerror(errno));
	}
	if (!compat) {
		WARN("%s is set to %d - set to 1 to ensure sched_yield works properly", pfile_name, compat);
	}
	fclose(pfile);
}

/*
// this echoes the definition that should be in sys/resource.h. This type is supported on 2.6.26 
// onwards but not documented in the man pages or available in the header files.
#ifndef RUSAGE_THREAD
#define RUSAGE_THREAD 1
#endif

double get_thread_exec_time() {
struct rusage usage;
CHECK_CALL(getrusage(RUSAGE_THREAD, &usage));
double usertime = (double)usage.ru_utime.tv_sec + (double)usage.ru_utime.tv_usec * 1e-6;
double systime = (double)usage.ru_stime.tv_sec + (double)usage.ru_stime.tv_usec * 1e-6;
return usage.ru_utime.tv_usec + usage.ru_stime.tv_usec;
//  return usertime + systime;
}
*/

void get_process_times(double *parent_time, double *children_time) {
	// the time for the children includes the forked and execed process, so we don't actually want 
	// that. The time for the parent includes all parent pthreads, so that is perfect
	struct tms t;
	CHECK_CALL(times(&t));
	double clk_tck = sysconf(_SC_CLK_TCK);
	*parent_time = (double)(t.tms_utime + t.tms_stime) / clk_tck;
	*children_time = (double)(t.tms_cutime + t.tms_cstime) / clk_tck;
	//	NOTICE("clock tick %.0f, process exec time %.3f, children exec time %.3f", 
	//				 clk_tck, *parent_time, *children_time);
}

double get_clock_res() {
	struct timespec ts;
	CHECK_CALL(syscall(__NR_clock_getres, CLOCK_REALTIME, &ts));
	return (double)ts.tv_sec + (double)ts.tv_nsec * 1e-9; 
}

double _utils_starting_time = 0;
double _utils_clk_tick = 0;

double init_starting_time() {
	_utils_clk_tick = (double)sysconf(_SC_CLK_TCK);  
	struct timespec ts = {.tv_sec = 0, .tv_nsec = 0};
	CHECK_CALL(syscall(__NR_clock_gettime, CLOCK_REALTIME, &ts));
	_utils_starting_time = (double)ts.tv_sec + (double)ts.tv_nsec * 1e-9;
	return _utils_clk_tick;
}

double get_current_time() {
	struct timespec ts;
	CHECK_CALL(syscall(__NR_clock_gettime, CLOCK_REALTIME, &ts));
	double clock_time = (double)ts.tv_sec + (double)ts.tv_nsec * 1e-9; 
	return clock_time - _utils_starting_time;
}

int get_rnd(int range, unsigned int *seed) {
	int rnd_index = ((double)rand_r(seed) * (double)range) / (double)RAND_MAX;
	if (rnd_index > range - 1 || rnd_index < 0) {
		FAIL("random out of range %d:  %d, aborting...\n", range, rnd_index);
	}
	return rnd_index;
}

void mx_init(pthread_mutex_t* mx, int attr) {
	pthread_mutexattr_t mx_attr;
	CHECK_PTHREAD_CALL(pthread_mutexattr_init(&mx_attr));
	CHECK_PTHREAD_CALL(pthread_mutexattr_settype(&mx_attr, attr));
	CHECK_PTHREAD_CALL(pthread_mutex_init(mx, &mx_attr));
	CHECK_PTHREAD_CALL(pthread_mutexattr_destroy(&mx_attr));
}

void set_process_affinity(int* cids, int num_cpus) {
	cpu_set_t cpu_set;
	CPU_ZERO(&cpu_set);
	// loop through all cpus and add to the cpu set
	for (int cid = 0; cid < num_cpus; cid++) CPU_SET(cids[cid], &cpu_set);
	CHECK_CALL(sched_setaffinity(getpid(), sizeof(cpu_set), &cpu_set));
}

int pin_thread(pid_t pid, int cid) {
	cpu_set_t cpu_set;
	CPU_ZERO(&cpu_set);
	CPU_SET(cid, &cpu_set);
	if (sched_setaffinity(pid, sizeof(cpu_set), &cpu_set) == -1) {
		if (errno == 3) WARN("%s, pid: %d", strerror(errno), pid);
		return -1;
	}
	return 0;
}

void check_is_pinned(pid_t ppid, pid_t pid, int cid) {
	int actual_cid = get_cid_from_procfs(ppid, pid);
	if (cid != actual_cid) {
		FAIL("couldn't pin pid %d to cpu %d: thread is on %d", pid, cid, actual_cid);
	}
}

int pin_thread_multi(pid_t pid, int* cids, int num_cpus) {
	cpu_set_t cpu_set;
	CPU_ZERO(&cpu_set);
	for (int i = 0; i < num_cpus; i++) CPU_SET(cids[i], &cpu_set);
	if (sched_setaffinity(pid, sizeof(cpu_set), &cpu_set) == -1) return -1;
	return 0;
}

typedef struct {
	int pid;
	char comm[16];
	char state;
	int ppid, pgrp, session, tty_nr, tpgid;
	unsigned int flags;
	unsigned long minflt, cminflt, majflt, cmajflt, utime, stime;
	long cutime, cstime, priority, nice, num_threads, itrealvalue;
	unsigned long long starttime;
	unsigned long vsize;
	long rss;
	unsigned long rsslim, startcode, endcode, startstack, kstkesp, kstkeip, signal, blocked, 
		sigignore, sigcatch, wchan, nswap, cnswap;
	int exit_signal, processor;
	unsigned int rt_priority, policy;
	unsigned long long delayacct_blkio_ticks;
	unsigned long guest_time;
	long cguest_time;
} task_status_t;

// this is very inefficient and should be avoided
int get_task_stats_from_procfs(char *procfs_name, task_status_t* tstat) {
	FILE* procfs_stats_file;
	if ((procfs_stats_file = fopen(procfs_name, "r")) == NULL) {
		return -1;
		//    FAIL("get_task_stats_from_procfs %s: %s", procfs_name, strerror(errno));
	}
	/*
	if (fscanf(procfs_stats_file, 
			   "%d %s %c %d %d %d %d %d %u %lu %lu %lu %lu %lu %lu %ld %ld %ld %ld %ld %ld %llu %lu "
			   "%ld %lu %lu %lu %lu %lu %lu %lu %lu %lu %lu %lu %lu %lu %d %d %u %u %llu %lu %ld\n",
			   &tstat->pid, tstat->comm, &tstat->state, &tstat->ppid, &tstat->pgrp, &tstat->session, 
			   &tstat->tty_nr, &tstat->tpgid, &tstat->flags, &tstat->minflt, &tstat->cminflt, 
			   &tstat->majflt, &tstat->cmajflt, &tstat->utime, &tstat->stime, &tstat->cutime, 
			   &tstat->cstime, &tstat->priority, &tstat->nice, &tstat->num_threads, 
			   &tstat->itrealvalue, &tstat->starttime, &tstat->vsize, &tstat->rss, &tstat->rsslim,
			   &tstat->startcode, &tstat->endcode, &tstat->startstack, &tstat->kstkesp, 
			   &tstat->kstkeip, &tstat->signal, &tstat->blocked, &tstat->sigignore, 
			   &tstat->sigcatch, &tstat->wchan, &tstat->nswap, &tstat->cnswap, &tstat->exit_signal, 
			   &tstat->processor, &tstat->rt_priority, &tstat->policy, &tstat->delayacct_blkio_ticks, 
			   &tstat->guest_time, &tstat->cguest_time) == EOF) {
	*/
	if (fscanf(procfs_stats_file, 
			   "%d %s %c %d %d %d %d %d %u %lu %lu %lu %lu %lu %lu ",
			   &tstat->pid, tstat->comm, &tstat->state, &tstat->ppid, &tstat->pgrp, &tstat->session, 
			   &tstat->tty_nr, &tstat->tpgid, &tstat->flags, &tstat->minflt, &tstat->cminflt, 
			   &tstat->majflt, &tstat->cmajflt, &tstat->utime, &tstat->stime) == EOF) {
		FAIL("[%s: %d] fscanf on %s: %s", __FUNCTION__, __LINE__, procfs_name, strerror(errno));
	}
	fclose(procfs_stats_file);
	return 0;
}

int is_task_running(pid_t pid) {
	FILE* pfile;
	char pfile_name[100];
	sprintf(pfile_name, "/proc/%d/status", pid); 
	if ((pfile = fopen(pfile_name, "r")) == NULL) {
		return 0;
		//		FAIL("open on %s: %s", pfile_name, strerror(errno));
	}
	char buf[500];
	if (fgets(buf, 500, pfile) == NULL) {
		return 0;
		//		FAIL("fgets on %s: %s", pfile_name, strerror(errno));
	}
	if (fgets(buf, 500, pfile) == NULL) {
		return 0;
		//		FAIL("fgets on %s: %s", pfile_name, strerror(errno));
	}
	//  NOTICE("pid %d state %s", pid, buf);
	fclose(pfile);
	if (!strstr(buf, "R (running)")) return 0;
	return 1;
}

int get_cid_from_procfs(pid_t ppid, pid_t pid) {
	task_status_t task_status;
	char procfs_fname[255];
	sprintf(procfs_fname, "/proc/%d/task/%d/stat", ppid, pid);
	if (get_task_stats_from_procfs(procfs_fname, &task_status) == -1) return -1;
	return task_status.processor;
}

// this is *very* expensive - not to be used in time critical loops
long long get_cpu_idle_ticks_slow(int cid) {
	FILE* pfile;
	char* pfile_name = "/proc/stat"; 
	if ((pfile = fopen(pfile_name, "r")) == NULL) {
		FAIL("fopen /proc/stat: %s", strerror(errno));
	}
	int cpu = -1;
	unsigned long utime, nice, stime, idle, iowait, irq, softirq, steal, guest, guest_nice; 
	for (int i = 0; i <= cid + 1; i++) {
		if (fscanf(pfile, "cpu%d %lu %lu %lu %lu %lu %lu %lu %lu %lu %lu\n", 
				   &cpu, &utime, &nice, &stime, &idle, &iowait, &irq, 
				   &softirq, &steal, &guest, &guest_nice) == EOF) {
			FAIL("fscanf on %s: %s", pfile_name, strerror(errno));
		}
		//printf("looking for cpu %d, found cpu %d\n", cid, cpu);
		if (cpu == cid) {
			fclose(pfile);
			return idle;
		}
	}
	fclose(pfile);
	return 0;
}

int _utils_per_cpu_stats_loaded = 1;

// this relies on having my module for getting cpu info out installed. Its cheap indeed.
long long get_cpu_idle_ticks(int cid, int fast) {
	if (!fast || !_utils_per_cpu_stats_loaded) return get_cpu_idle_ticks_slow(cid);
	FILE* pfile = NULL;
	char pfile_name[255];
	sprintf(pfile_name, "/proc/per_cpu_stats/cpu%d", cid);
	//  printf("pfile_name = %s\n", pfile_name);
	if ((pfile = fopen(pfile_name, "r")) == NULL) {
		WARN("Cannot open %s. Using slow method to get CPU idle information. "
			 "Load the per_cpu_stats module for the fast method.", pfile_name);
		_utils_per_cpu_stats_loaded = 0;
		return get_cpu_idle_ticks_slow(cid);
	}
	char buff[101];
	if (fgets(buff, 100, pfile) == NULL) {
		FAIL("fscanf on %s: %s\n", pfile_name, strerror(errno));
	}
	fclose(pfile);
	return atoll(buff);
}

  
int get_task_stats_from_proc(pid_t pid, pid_t tid, 
							 double *utime, double *stime, double *rtime) {
	char procfs_fname[255];
	sprintf(procfs_fname, "/proc/%d/task/%d/stat", pid, tid);
	task_status_t task_status;
	if (get_task_stats_from_procfs(procfs_fname, &task_status) == -1) return -1;
	*utime = task_status.utime / _utils_clk_tick;
	*stime = task_status.stime / _utils_clk_tick;
	*rtime = get_current_time();
	return 0;
}

int get_hyperthreads_per_core() {
	FILE* f = fopen("/proc/cpuinfo", "r");
	if (!f) {
		FAIL("could not open %s", "/proc/cpuinfo");
	}
	char buf[512];
	for (int i = 0; i < 9; i++) {
		if (!fgets(buf, 511, f)) {
			FAIL("could not read %s", "/proc/cpuinfo");
		}
		//    printf(buf);
	}
	int siblings, cores;
	fscanf(f, "siblings\t: %d\n", &siblings);
	fgets(buf, 511, f);
	//  printf(buf);
	fscanf(f, "cpu cores\t: %d\n", &cores);
	//  printf("siblings %d cores %d\n", siblings, cores);
	fclose(f);
	return siblings / cores;
}
